package kibahed.uni.mdp.rlalgorithms;

import java.text.DecimalFormat;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;

import kibahed.uni.mdp.representation.Action;
import kibahed.uni.mdp.representation.MDPField;
import kibahed.uni.mdp.representation.Position;
import kibahed.uni.mdp.representation.State;

public class PAlgorithm extends ARLalgorithm {

	private Map<State, Map<Action, Integer>> actionCount = new HashMap<State, Map<Action, Integer>>();
	private Map<State, Map<Action, Integer>> goodTries = new HashMap<State, Map<Action, Integer>>();
	private Map<State, Map<Action, Map<State, Integer>>> walks = new HashMap<State, Map<Action, Map<State, Integer>>>();
	private Map<State, Map<Action, Map<State, Double>>> pHats = new HashMap<State, Map<Action, Map<State, Double>>>();

	private int N = 0;

	private Map<State, Map<Action, Double>> T_s_a = new HashMap<State, Map<Action, Double>>();
	private Map<State, Double> T_s;

	private int stopTevaluation = 10;
	private double continueWalkingPercentage = 1.1;

	public PAlgorithm(MDPField mdpField) {
		super(mdpField);
		initT();
		initMaps();

	}

	private void initMaps() {
		walks = new HashMap<State, Map<Action, Map<State, Integer>>>();
		for (State s1 : mdpField.getStateList()) {
			Map<Action, Map<State, Integer>> m1 = new HashMap<Action, Map<State, Integer>>();
			for (Action a : s1.getActions()) {
				Map<State, Integer> m2 = new HashMap<State, Integer>();
				for (State s2 : s1.getNeighbours()) {
					m2.put(s2, 0);
				}
				m1.put(a, m2);
			}
			walks.put(s1, m1);
		}

		actionCount = new HashMap<State, Map<Action, Integer>>();
		for (State s : mdpField.getStateList()) {
			Map<Action, Integer> m = new HashMap<Action, Integer>();
			for (Action a : s.getActions()) {
				m.put(a, 1);
			}
			actionCount.put(s, m);
		}
	}

	private void initT() {
		T_s_a = new HashMap<State, Map<Action, Double>>();
		T_s = new HashMap<State, Double>();
		for (State s : mdpField.getStateList()) {
			Map<Action, Double> m = new HashMap<Action, Double>();
			for (Action a : s.getActions()) {
				m.put(a, 0.);
			}
			T_s_a.put(s, m);
			T_s.put(s, 0.);
		}
	}

	@Override
	public void calculate() {
		Position pos = mdpField.getRandomPosition();
		// how many steps should be done?
		int steps = 1000;
		do {

			Map<State, Map<Action, Integer>> oldActionCount = new HashMap<State, Map<Action, Integer>>(
					actionCount);
			for (State s : actionCount.keySet()) {
				Map<Action, Integer> help = new HashMap<Action, Integer>(
						actionCount.get(s));
				oldActionCount.put(s, help);
			}
			do {
				walkAround(pos, steps);
			} while (!walkedEnough(oldActionCount));

			// walkAround(pos, steps);

		} while (updateT(stopTevaluation));

		printTValues();
		print();
		printParamters();
		printPolicy();
	}

	private boolean walkedEnough(Map<State, Map<Action, Integer>> oldActionCount) {
		// System.out.println("walkAround");
		for (State s : oldActionCount.keySet()) {
			for (Action a : oldActionCount.get(s).keySet()) {
				// System.out.println("old AC = "+oldActionCount.get(s).get(a)+" new AC = "+actionCount.get(s).get(a));
				if (actionCount.get(s).get(a) > continueWalkingPercentage
						* oldActionCount.get(s).get(a)) {
					// System.out.println("made double steps");
					return true;
				}
			}
		}
		return false;
	}

	public void doWalking(Position start, int steps) {
		walkAround(start, steps);
		print();
	}

	private void walkAround(Position startPosition, int steps) {
		// System.out.println("Steps left = " + steps);
		State start = mdpField.getState(startPosition);
		Random r = new Random();
		int count = 0;
		while (count < steps) {
			Action a;
			if (policy.isEmpty()) {
				int index = r.nextInt(start.getActions().size());
				a = start.getActions().get(index);
			} else {
				a = policy.get(start);
			}
			incrementActionCount(start, a);
			State nextState = new State(start.takeAction(a));
			adaptPhats(start, a, nextState);
			start = nextState;
			count++;
			if (start.isAbsorbing()) {
				start = getStateWithMinVisits();
				walkAround(start.getPosition(), steps - count);
				break;
			}
		}

	}

	private State getStateWithMinVisits() {
		State selectedState = null;
		int min = Integer.MAX_VALUE;
		for(State s : actionCount.keySet()) {
			int sum = sumUpStepsPerState(s);
			if(sum < min) {
				min = sum;
				selectedState = s;
			}
		}
		return selectedState;
	}

	private boolean updateT(int minIterations) {

		Map<State, Double> oldT;
		int it = 0;
		do {
			oldT = new HashMap<State, Double>(T_s);

			for (State s : T_s.keySet()) {

				if (s.isAbsorbing()) {

					for (Action a : T_s_a.get(s).keySet()) {
						Map<Action, Double> m = T_s_a.get(s);
						m.put(a, 0.);
						T_s_a.put(s, m);
					}

				} else {

					for (Action a : T_s_a.get(s).keySet()) {

						// here is the state-action pair !!
						// create sum over all s'

						double sum = 0;
						double nsa = 1;
						if (actionCount.containsKey(s)
								&& actionCount.get(s).containsKey(a)) {
							nsa = (double) actionCount.get(s).get(a);
						}

						if (N == 0)
							N++;
						double confidence = Math.sqrt(Math.log((double) N)
								/ nsa);

						if (pHats.containsKey(s) && pHats.get(s).containsKey(a)) {

							for (State nextState : pHats.get(s).get(a).keySet()) {

								double pHat = pHats.get(s).get(a)
										.get(nextState);

								double t = T_s.get(nextState);

								sum += pHat * t;

							}
						}

						double additionValue = sum - confidence;

						Map<Action, Double> m = T_s_a.get(s);
						m.put(a, 1 + additionValue);
						T_s_a.put(s, m);
					}

				}

				T_s.put(s, getMinTsa(s));
				// System.out.println("T_T_S = " + T_s.get(s));
			}

			it++;
		} while (!convergationReached(oldT, 0.0001));

		// stop if only needed less iterations then minIterations to convergate
		if (it > minIterations) {
			return true;
		} else {
			return false;
		}

	}

	private double getMinTsa(State state) {
		double min = Double.MAX_VALUE;
		Action policyAction = null;
		for (Action a : T_s_a.get(state).keySet()) {
			if (T_s_a.get(state).get(a) < min) {
				min = T_s_a.get(state).get(a);
				policyAction = a;
			}
		}

		if (state.isAbsorbing()) {
			policy.put(state, new Action(0, "stay"));
		} else {
			policy.put(state, policyAction);
		}
		return min;
	}

	private double sumUpTs(State state) {
		double sum = 0;
		if (T_s_a.containsKey(state)) {
			for (Action a : T_s_a.get(state).keySet()) {
				sum += T_s_a.get(state).get(a);
			}
		}
		return sum;
	}

	private boolean convergationReached(Map<State, Double> oldT, double epsilon) {
		int count = 0;
		for (State s : T_s.keySet()) {

			if ((Math.abs(T_s.get(s) - oldT.get(s)) < epsilon && T_s.get(s) != 1)) {
				count++;
			}
		}
		if (count == oldT.keySet().size()) {
			return true;
		}
		return false;
	}

	@Override
	public void calculate(Position startPosition) {

	}

	private void incrementActionCount(State state, Action a) {
		if (actionCount.containsKey(state)) {
			if (actionCount.get(state).containsKey(a)) {
				actionCount.get(state)
						.put(a, actionCount.get(state).get(a) + 1);
			} else {
				actionCount.get(state).put(a, 1);
			}
		} else {
			Map<Action, Integer> map = new HashMap<Action, Integer>();
			map.put(a, 1);
			actionCount.put(state, map);
		}
		N++;
	}

	private void adaptPhats(State start, Action a, State goal) {

		// check if try was good
		if (start.getAimStateForAction(a).equals(goal)) {
			if (goodTries.containsKey(start)) {
				if (goodTries.get(start).containsKey(a)) {
					goodTries.get(start)
							.put(a, goodTries.get(start).get(a) + 1);
				} else {
					goodTries.get(start).put(a, 1);
				}
			} else {
				Map<Action, Integer> map = new HashMap<Action, Integer>();
				map.put(a, 1);
				goodTries.put(start, map);
			}

		}

		walks.get(start).get(a)
				.put(goal, walks.get(start).get(a).get(goal) + 1);

		// evaluate new pHat for state start, action a and state goal;

		int tries = getTries(start, a);

		int goodOnes = getGoodOnes(start, a);

		double p = goodOnes / (double) tries;

		if (pHats.containsKey(start)) {
			if (pHats.get(start).containsKey(a)) {
				pHats.get(start).get(a).put(goal, p);

			} else {
				Map<State, Double> m = new HashMap<State, Double>();
				m.put(goal, p);
				pHats.get(start).put(a, m);
			}
		} else {
			Map<Action, Map<State, Double>> m1 = new HashMap<Action, Map<State, Double>>();
			Map<State, Double> m2 = new HashMap<State, Double>();
			m2.put(goal, p);
			m1.put(a, m2);
			pHats.put(start, m1);
		}
		calculatePHats(start, a);
	}

	private void calculatePHats(State state, Action action) {
		for (State s : walks.get(state).get(action).keySet()) {
			pHats.get(state)
					.get(action)
					.put(s,
							walks.get(state).get(action).get(s)
									/ (double) actionCount.get(state).get(
											action));
		}
	}

	private void evalPHats(State start, Action action, State goal) {

		for (Action a : pHats.get(start).keySet()) {
			int tries = getTries(start, a);
			int goodOnes = getGoodOnes(start, a);
			double p = goodOnes / (double) tries;
			pHats.get(start).get(a).put(goal, p);
		}

		double sum = 0;
		for (Action a : pHats.get(start).keySet()) {
			for (State s : pHats.get(start).get(a).keySet()) {
				sum += pHats.get(start).get(a).get(s);
			}
			for (State s : pHats.get(start).get(a).keySet()) {

				pHats.get(start).get(a)
						.put(s, pHats.get(start).get(a).get(s) / sum);
			}
		}
	}

	private int getGoodOnes(State start, Action a) {
		int goodOnes = 0;
		if (goodTries.containsKey(start)) {
			if (goodTries.get(start).containsKey(a)) {
				goodOnes = goodTries.get(start).get(a);
			}
		}
		return goodOnes;
	}

	private int getTries(State start, Action a) {
		int tries = 0;
		if (actionCount.containsKey(start)) {
			if (actionCount.get(start).containsKey(a)) {
				tries = actionCount.get(start).get(a);
			}
		}
		return tries;

	}

	private void print() {

		System.out
				.println("------------------------------------------------------------- \n"
						+ "NUMBER OF ACTIONS IN EACH STATE - P ALGORITHM \n"
						+ "--------------------------------------------------------------\n");
		DecimalFormat df = new DecimalFormat();
		df.setMaximumFractionDigits(0);
		for (int i = 0; i < mdpField.getStateList().size(); i++) {
			State s = mdpField.getStateList().get(i);
			System.out.print(" | [" + s.getPosition().getRow() + ","
					+ s.getPosition().getCollum() + "] # = "
					+ df.format(sumUpStepsPerState(s)));
			if ((i + 1) % (mdpField.getCollums()) == 0 && i != 0) {
				System.out.println("\n");
			}
		}
		System.out.print("|");

		System.out.println("number of steps in total = " + df.format(N));
	}

	private void printPHats() {
		for (State s : pHats.keySet()) {
			System.out.println("\n\nState " + s.printPosition());
			for (Action a : pHats.get(s).keySet()) {
				System.out.println("\nAction " + a.getActionName() + "\n");
				double sum = 0;
				for (State n : pHats.get(s).get(a).keySet()) {
					sum += pHats.get(s).get(a).get(n);
					System.out.print(" |  " + pHats.get(s).get(a).get(n));
				}
				System.out.print(" | sum = " + sum);
			}
		}
	}

	private int sumUpStepsPerState(State state) {
		int count = 0;
		if (!state.isAbsorbing()) {
			if (actionCount.containsKey(state)) {
				for (Action a : actionCount.get(state).keySet()) {
					count += actionCount.get(state).get(a);
				}
			}
		}
		return count;
	}

	private void printTValues() {

		System.out
				.println("------------------------------------------------------------- \n"
						+ "RESULT FOR P ALGORITHM \n"
						+ "--------------------------------------------------------------\n");

		DecimalFormat df = new DecimalFormat();
		df.setMaximumFractionDigits(2);
		for (int i = 0; i < mdpField.getStateList().size(); i++) {
			State s = mdpField.getStateList().get(i);
			System.out.print(" | [" + s.getPosition().getRow() + ","
					+ s.getPosition().getCollum() + "] T = "
					+ df.format(T_s.get(s)));
			if ((i + 1) % (mdpField.getCollums()) == 0 && i != 0) {
				System.out.println("\n");
			}
		}
		System.out.print("|");
	}

	private void printParamters() {

		DecimalFormat df = new DecimalFormat();
		df.setMaximumFractionDigits(0);
		System.out.println("\nParameter for PAlgorithm:\n");
		System.out.println("stopTevaluation = " + stopTevaluation);
		System.out.println("continueWalkingPercentage = "
				+ continueWalkingPercentage);

	}

	private void printPolicy() {

		System.out
				.println("------------------------------------------------------------- \n"
						+ "POLICY \n"
						+ "--------------------------------------------------------------\n");

		DecimalFormat df = new DecimalFormat();
		df.setMaximumFractionDigits(2);
		for (int i = 0; i < mdpField.getStateList().size(); i++) {
			State s = mdpField.getStateList().get(i);
			System.out.print(" | [" + s.getPosition().getRow() + ","
					+ s.getPosition().getCollum() + "] : Action = "
					+ policy.get(s).getActionName());
			if ((i + 1) % (mdpField.getCollums()) == 0 && i != 0) {
				System.out.println("\n");
			}
		}
		System.out.print("|");
	}
}
